# 数据集加载，这里加载的数据都是公共的图神经网络数据集，直接下载然后加载就可以了

from torch_geometric.loader import DataLoader
from torch_geometric.datasets import TUDataset

# 加载TUDataset中的数据集，如果没有则下载到root指定的文件中
dataset = TUDataset(
    root='../../datasets/ENZYMES',
    name='ENZYMES',
    use_node_attr=True,
)
# 建立数据集加载器，每次加载 batch_size 个数据到内存中
loader = DataLoader(
    # 要加载的数据集
    dataset=dataset,
    # ENZYMES包含600个图,每次加载32个
    batch_size=32,
    # 每次加入进来之后是否随机打乱数据（可以增加模型的泛化性）
    shuffle=True
)
# 遍历所有的 batch
for batch in loader:
    print(batch)
    print(batch.num_graphs)
